# Copyright 2009, Richard Monk <richardmonk@gmail.com>
#
# This file is part of pipestat.
#
# pipestat is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# pipestat is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with pipestat.  If not, see <http://www.gnu.org/licenses/>.
"""
You can put this script in a pipe and get output to stderr regarding the data
transfer like so:
Transferring   38.60 KB/sec , overall   35.77 KB/sec, total    3.10 MB

Very useful for watching network transfers, or basic testing of drive
speeds or program output rate
"""

__docformat__ = 'restructuredtext'
# Major.minor.patch. Prereleases use patch like .01, never .1
__version__ = '0.4.2'
__license__ = 'GPLv3+'
__author__ = "Richard Monk"


import fcntl
import optparse
import os
import platform
import select
import sys
import time
try:
    import hashlib
    oldhash = False
except ImportError:
# Needed for old python instances without hashlib
    import sha
    import md5
    oldhash = True

# Make read requests to stdin nonblocking
fcntl.fcntl(0, fcntl.F_SETFL, os.O_NONBLOCK)


def prettyPrint(value, trailer=""):
    """
    PrettyPrint just coverts a value to its KB/MB/GB format
    """
    rates = [("GB", 1024 ** 3), ("MB", 1024 ** 2), ("KB", 1024)]
    divider = 1
    rateunit = "B"

    for unit, size in rates:
        if value > size:
            divider = size
            rateunit = unit
            break

    return "%7.2f %s%s" % ((float(value) / divider), rateunit, trailer)


def main():
    """
    Main entry point.
    """
    # Exit if this is windows
    if "".join(platform.win32_ver()):
        sys.stderr.write("This tool does not work on Windows. Exiting ...\n")
        raise SystemExit(1)

    # Initalize the sha1 and md5 hashes
    if oldhash:
        md5hash = md5.new()
        sha1hash = sha1.new()
    else:
        md5hash = hashlib.md5()
        sha1hash = hashlib.sha1()

    # Setup the parser
    parser = optparse.OptionParser(version="%prog " + __version__)
    (options, args) = parser.parse_args()


    starttime = time.time()
    last_data_transferred = 0
    data_transferred = 0
    currtime = time.time()
    fileno = sys.stdin.fileno()
    readset = [sys.stdin]
    persec = total = "bytes"
    buf = ''

    try:
        # Run until the stdin dies
        while not sys.stdin.closed:
            try:
                # Check every second for updates, if none, just keep going
                # (to keep the stats updated)
                inready, outready, eready = select.select(readset, [], [], 1)
                if inready != []:
                    buf = sys.stdin.read(32768)
                    if not buf:
                        # Bail out, buffer is closed
                        raise SystemExit(1)

                # Store time and amount of data transferred for later display
                newtime = time.time()
                data_transferred += len(buf)
                if inready != []:
                    sys.stdout.write(buf)
                    md5hash.update(buf)
                    sha1hash.update(buf)

                timediff = newtime - currtime
                # only update every quarter second max
                if timediff >= .25:
                    bytes_per_sec = (
                        data_transferred - last_data_transferred) / timediff

                    data_per_sec = prettyPrint(bytes_per_sec, "/sec")
                    overall_per_sec = (
                        data_transferred / (newtime - starttime))
                    overall_per_sec = prettyPrint(overall_per_sec, "/sec")

                    sys.stderr.write(("Transferring %s , overall %s, total "
                        "%s\r" % (data_per_sec, overall_per_sec,
                        prettyPrint(data_transferred))))
                    sys.stderr.flush()
                    last_data_transferred = data_transferred
                    currtime = newtime
                    buf = ""
            except KeyboardInterrupt, ki:
                raise SystemExit(0)
            except Exception, ex:
                raise SystemExit(2)
    except SystemExit, se:
        # errors mean one end is dead.  Close out.
        sys.stderr.write("\nWrote %s bytes in %s seconds\n" % (
            prettyPrint(data_transferred), int(currtime - starttime)))
        sys.stderr.write("md5 : %s\n" % (md5hash.hexdigest()))
        sys.stderr.write("sha1: %s\n" % (sha1hash.hexdigest()))
        sys.stderr.flush()
        raise se

if __name__ == "__main__":
    main()
